import operator
from typing import List, Dict, TypedDict, Annotated

from dotenv import load_dotenv
from langchain_core.messages import HumanMessage
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, END

load_dotenv()
import os


# 1. 定义全局状态
class GameState(TypedDict):
    # 记录所有人的发言
    history: Annotated[List[str], operator.add]
    # 记录每个人的词 (上帝视角)
    player_words: Dict[str, str]
    # 记录投票结果
    votes: Dict[str, str]
    # 游戏结果
    winner: str


# 初始化模型
llm = ChatOpenAI(model=os.getenv("LONGMAO_MODEL"),
                 base_url=os.getenv("LONGMAO_BASE_URL"),
                 api_key=os.getenv("LONGMAO_API_KEY"))


# ----------------- 节点逻辑 -----------------

# 【上帝节点】：负责发牌
def assign_roles_node(state: GameState):
    print("--- 🎲 游戏开始，上帝正在发牌 ---")
    # 假设词库：苹果 vs 梨
    # 0是平民词，1是卧底词
    words = ["苹果", "梨"]

    # 随机分配：A, B, C 三个玩家
    # 设定：A和C是平民(苹果)，B是卧底(梨)
    allocation = {
        "Player_A": words[0],
        "Player_B": words[1],
        "Player_C": words[0]
    }
    print(f"（上帝视角：分配结果 {allocation}）")
    return {"player_words": allocation, "history": []}


# 【玩家节点工厂】：生成玩家的发言逻辑
def make_player_node(player_name: str):
    def player_node(state: GameState):
        my_word = state["player_words"][player_name]
        history_text = "\n".join(state["history"])

        # 构造 Prompt：私有词汇 + 公共历史
        prompt = f"""
        你正在玩【谁是卧底】游戏。共有3人参与。
        你的身份是：{player_name}。
        你的关键词是：【{my_word}】。

        请用一句话描述你的词，不要直接说出这个词，也不要太明显，否则会被投出去。
        同时你要观察别人的发言，看谁像卧底。

        之前的发言记录：
        {history_text}

        请输出你的描述：
        """

        response = llm.invoke([HumanMessage(content=prompt)])
        content = f"{player_name}: {response.content}"
        print(f"🗣️ {content}")

        return {"history": [content]}

    return player_node


# 【投票节点工厂】：生成玩家的投票逻辑
def make_vote_node(player_name: str):
    def vote_node(state: GameState):
        history_text = "\n".join(state["history"])
        my_word = state["player_words"][player_name]

        prompt = f"""
        我是 {player_name}，我的词是【{my_word}】。
        根据以下历史发言，请投票选出谁是卧底（词和大家不一样的那个）。
        只能输出玩家名字（Player_A, Player_B, 或 Player_C），不要输出其他内容。

        历史发言：
        {history_text}
        """

        response = llm.invoke([HumanMessage(content=prompt)])
        vote_target = response.content.strip()
        print(f"🗳️ {player_name} 投给了 -> {vote_target}")

        # 更新投票状态
        return {"votes": {player_name: vote_target}}

    return vote_node


# 【判决节点】：计算票数，决定胜负
def judge_node(state: GameState):
    print("--- ⚖️ 投票结束，上帝正在计票 ---")
    votes = state["votes"]
    # 统计票数
    vote_counts = {}
    for voter, target in votes.items():
        vote_counts[target] = vote_counts.get(target, 0) + 1

    # 找到得票最多的人
    eliminated = max(vote_counts, key=vote_counts.get)
    print(f"❌ {eliminated} 被高票淘汰！")

    # 判断胜负
    player_words = state["player_words"]

    # 逻辑：如果卧底(Player_B)被淘汰，平民赢；否则卧底赢
    # 真实场景这里应该动态判断谁拿了哪个词，这里为了演示简化写死逻辑
    spy_name = [name for name, word in player_words.items() if word == "梨"][0]

    if eliminated == spy_name:
        result = "平民胜利！卧底被抓住了！"
    else:
        result = f"卧底胜利！平民 {eliminated} 被误杀！"

    print(f"🏆 结果：{result}")
    return {"winner": result}


# ----------------- 构建图 (LangGraph) -----------------

workflow = StateGraph(GameState)

# 1. 添加节点
workflow.add_node("God_Assign", assign_roles_node)
workflow.add_node("Player_A_Speak", make_player_node("Player_A"))
workflow.add_node("Player_B_Speak", make_player_node("Player_B"))
workflow.add_node("Player_C_Speak", make_player_node("Player_C"))
workflow.add_node("Player_A_Vote", make_vote_node("Player_A"))
workflow.add_node("Player_B_Vote", make_vote_node("Player_B"))
workflow.add_node("Player_C_Vote", make_vote_node("Player_C"))
workflow.add_node("God_Judge", judge_node)

# 2. 定义边 (流程控制)
# 发牌 -> A说 -> B说 -> C说 -> A投 -> B投 -> C投 -> 判决
workflow.set_entry_point("God_Assign")
workflow.add_edge("God_Assign", "Player_A_Speak")
workflow.add_edge("Player_A_Speak", "Player_B_Speak")
workflow.add_edge("Player_B_Speak", "Player_C_Speak")
workflow.add_edge("Player_C_Speak", "Player_A_Vote")  # 发言结束，开始投票
workflow.add_edge("Player_A_Vote", "Player_B_Vote")
workflow.add_edge("Player_B_Vote", "Player_C_Vote")
workflow.add_edge("Player_C_Vote", "God_Judge")
workflow.add_edge("God_Judge", END)

# 3. 编译应用
app = workflow.compile()

# ----------------- 运行 -----------------
if __name__ == "__main__":
    app.invoke({"history": [], "votes": {}, "player_words": {}, "winner": ""})
